# The `ForcedPhotometry` class performs forced photometry on astronomical images, supporting both
# global and local background estimation methods.
from loguru import logger
import numpy as np
from astropy.io import fits
from astropy.wcs import WCS
from astropy.coordinates import SkyCoord
import astropy.units as u
from photutils import SkyCircularAperture, aperture_photometry, CircularAnnulus
from photutils.background import Background2D, MedianBackground

from tools.utils import MatrixPartitioner, ra_dec_to_xy_and_size, read_header_keys

class ForcedPhotometry:
    """
    强制测光类，支持全局背景和局部背景的计算。

    修正说明：
    - 数据是ADU
    - 所有计算基于ADU数，噪声模型和星等公式已修正
    """

    def __init__(self, fitsfile, zero_point):
        self.fitsfile = fitsfile
        self.zero_point = zero_point  # 零点单位需为 mag/(ADU/s)
        self.data, self.wcs, self.exptime, self.gain, self.readnoise = self._load_fits()

    def _load_fits(self):
        """加载FITS文件，读书是ADU"""
        with fits.open(self.fitsfile) as hdul:
            raw_data = hdul[0].data.astype(float)
            header = hdul[0].header
            wcs = WCS(header)
            exptime = header.get('EXPOSURE', 1.0)
            gain  = header.get('GAIN', 1.0)
            readnoise = header.get('RDNOISE', 0.0)
            data = raw_data  #单位为ADU
        logger.debug(f"Loaded FITS. Exposure={exptime}s, Gain={gain}, Readnoise={readnoise}e⁻")
        return data, wcs, exptime, gain, readnoise

    def _calculate_global_background(self):
        """计算全局背景（基于ADU数）"""
        bkg_estimator = MedianBackground()
        bkg = Background2D(self.data, (128, 128), filter_size=(3, 3), bkg_estimator=bkg_estimator)
        logger.debug(f"Global Bkg: median={np.median(bkg.background):.1f}, RMS={np.median(bkg.background_rms):.1f}")
        return bkg.background, bkg.background_rms

    def _calculate_local_background(self, x, y, aperture_radius_pix):
        """计算局部背景（基于ADU数）"""
        annulus = CircularAnnulus((x, y), r_in=aperture_radius_pix*1.5, r_out=aperture_radius_pix*2.0)
        annulus_mask = annulus.to_mask(method='center')
        annulus_data = annulus_mask.multiply(self.data)
        annulus_data_1d = annulus_data[np.isfinite(annulus_data) & (annulus_data > 0)]
        
        if annulus_data_1d.size == 0:
            logger.warning("Local annulus invalid. Using global median.")
            local_background = np.median(self.data)
            local_background_rms = np.std(self.data)
        else:
            local_background = np.median(annulus_data_1d)
            local_background_rms = np.std(annulus_data_1d)
        logger.debug(f"Local Bkg: median={local_background:.1f}ADU, RMS={local_background_rms:.1f}ADU")
        return local_background, local_background_rms

    def perform_photometry(self, ra, dec, aperture_radius_arcsec, use_local_background=False):
        """执行测光（所有计算基于电子数）"""
        
        # 坐标转换
        coords = SkyCoord(ra*u.deg, dec*u.deg, frame='icrs')
        apertures = SkyCircularAperture(coords, r=aperture_radius_arcsec*u.arcsec)
        pixel_apertures = apertures.to_pixel(self.wcs)

        # 计算孔径像素数
        tid = self.fitsfile.split('/')[-1][:2].lower()
        pixel_scale = 0.286 if tid == 'mr' else 0.429  # 根据文件名硬编码像素比例
        aperture_radius_pix = aperture_radius_arcsec / pixel_scale
        n_pixels = np.pi * (aperture_radius_pix**2)

        # 测光总通量（ADU数）
        phot_table = aperture_photometry(self.data, pixel_apertures)
        total_flux = phot_table['aperture_sum'][0]
        logger.debug(f"Total Flux: {total_flux:.1f}ADU")

        # 背景计算
        x, y = int(phot_table['xcenter'][0].value), int(phot_table['ycenter'][0].value)
        if use_local_background:
            bkg, bkg_rms = self._calculate_local_background(x, y, aperture_radius_pix)
            bkg_value = n_pixels * bkg
            bkg_noise = np.sqrt(n_pixels) * bkg_rms
        else:
            global_bkg, global_bkg_rms = self._calculate_global_background()
            bkg_value = n_pixels * global_bkg[y, x]
            bkg_noise = np.sqrt(n_pixels) * global_bkg_rms[y, x]

        # 净通量与噪声模型
        net_flux = total_flux - bkg_value
        poisson_variance = net_flux  # Poisson噪声方差
        limit_use_bkg_rms=bkg_rms if use_local_background else global_bkg_rms[y,x]
        bkg_variance = (bkg_rms**2) * n_pixels if use_local_background else (global_bkg_rms[y,x]**2) * n_pixels
        readnoise_variance = n_pixels * ((self.readnoise/self.gain)**2)
        variance = poisson_variance/self.gain + bkg_variance + readnoise_variance
        snr = net_flux / np.sqrt(variance) if variance > 0 else 0

        # 星等计算
        magnitude = self.zero_point - 2.5 * np.log10(net_flux / self.exptime)
        
        limiting_mag = self.zero_point - 2.5 * np.log10(
            3 * (limit_use_bkg_rms / self.exptime) * np.sqrt(2.226 * (5 ** 2))
        )

        return {
            "total_flux": total_flux,
            "net_flux": net_flux,
            "snr": snr,
            "magnitude": magnitude,
            "limiting_mag": limiting_mag
        }
    def generate_mag_growth_curve(self, ra, dec, aperture_range=(1.0, 5.0), steps=50, output_path=None):
        """
        生成星等生长曲线（不同孔径下的星等变化）
        :param aperture_range: 孔径范围（角秒），格式(start, end)
        :param steps: 测试点数
        :param output_path: 图片保存路径
        """
        import matplotlib.pyplot as plt
        from scipy.signal import savgol_filter
        
        apertures = np.linspace(aperture_range[0], aperture_range[1], steps)
        magnitudes = []
        snr_values = []
        
        # 遍历不同孔径
        for ap in apertures:
            result = self.perform_photometry(ra, dec, ap)
            magnitudes.append(result["magnitude"])
            snr_values.append(result["snr"])
            logger.debug(f"孔径 {ap:.1f}\": 星等 = {result['magnitude']:.2f}")

        # 数据平滑处理
        smoothed = savgol_filter(magnitudes, window_length=5, polyorder=2)
        
        # 寻找稳定点（变化率小于0.05 mag/arcsec）
        gradient = np.gradient(smoothed)
        stable_points = np.where(np.abs(gradient) < 0.05)[0]
        optimal_ap = apertures[stable_points[0]] if len(stable_points) > 0 else apertures[-1]

        # 绘制曲线
        plt.figure(figsize=(12, 6))
        
        # 星等曲线
        plt.subplot(211)
        plt.plot(apertures, magnitudes, 'bo-', label='origin data')
        plt.plot(apertures, smoothed, 'r--', lw=2, label='smooth curve')
        plt.axvline(optimal_ap, color='g', linestyle='--', label=f'statble app ({optimal_ap:.1f}")')
        plt.gca().invert_yaxis()
        plt.ylabel('mag (AB)')
        plt.title('mag growth curve')
        plt.legend()
        plt.grid(True)
        
        # SNR曲线
        plt.subplot(212)
        plt.plot(apertures, snr_values, 'mD--', markersize=6)
        plt.xlabel('Aperture radius (arcseconds)')
        plt.ylabel('SNR')
        plt.grid(True)
        
        plt.tight_layout()

        if output_path:
            plt.savefig(output_path, dpi=150)
            logger.info(f"星等生长曲线已保存至 {output_path}")
        else:
            plt.show()
        
        return {
            'apertures': apertures,
            'magnitudes': magnitudes,
            'optimal_aperture': optimal_ap,
            'snr_values': snr_values
        }

    def generate_light_curve(self, data_dir, ra, dec, aperture_radius_arcsec, output_path=None):
        """
        生成目标的光变曲线
        :param data_dir: 包含时序FITS文件的目录
        :param ra: 目标赤经（度）
        :param dec: 目标赤纬（度）
        :param aperture_radius_arcsec: 测光孔径半径（角秒）
        :param output_path: 图片保存路径（None则显示在窗口）
        """
        from pathlib import Path
        import matplotlib.pyplot as plt
        from astropy.time import Time
        
        files = sorted(Path(data_dir).glob("*.fits"))
        jd_times = []
        magnitudes = []
        snr_values = []
        
        for f in files:
            try:
                # 创建新的测光实例处理每个文件
                phot = ForcedPhotometry(str(f), self.zero_point)
                
                # 从文件头获取精确时间
                with fits.open(str(f)) as hdul:
                    header = hdul[0].header
                    date_obs = header.get('DATE-OBS', '2000-01-01T00:00:00.000')  # 假设头信息包含DATE-OBS
                t = Time(date_obs, format='isot').jd  # 转换为儒略日
                
                # 执行测光
                result = phot.perform_photometry(ra, dec, aperture_radius_arcsec)
                
                # 记录有效数据
                if result["magnitude"] is not None:
                    jd_times.append(t)
                    magnitudes.append(result["magnitude"])
                    snr_values.append(result["snr"])
                    
            except Exception as e:
                logger.error(f"处理文件 {f.name} 失败: {str(e)}")
                continue

        # 绘制光变曲线
        plt.figure(figsize=(12, 6))
        
        # 星等面板
        plt.subplot(211)
        plt.errorbar(jd_times, magnitudes, fmt='o', markersize=5, 
                    capsize=3, ecolor='gray', alpha=0.7)
        plt.gca().invert_yaxis()
        plt.ylabel('Magnitude (AB)')
        plt.title(f'Light Curve (Aperture: {aperture_radius_arcsec}" )')
        
        # SNR面板
        plt.subplot(212)
        plt.plot(jd_times, snr_values, 'g^--', markersize=5, alpha=0.7)
        plt.ylabel('SNR')
        plt.xlabel('Julian Date')
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=150)
            logger.info(f"光变曲线已保存至 {output_path}")
        else:
            plt.show()
        
        return jd_times, magnitudes, snr_values

    def generate_snr_aperture_curve(self, ra, dec, aperture_range=(1.0, 5.0), steps=10, output_path=None):
        """
        生成信噪比-孔径曲线
        :param aperture_range: 孔径范围（角秒），格式(start, end)
        :param steps: 测试点数
        :param output_path: 图片保存路径
        """
        import matplotlib.pyplot as plt
        
        apertures = np.linspace(aperture_range[0], aperture_range[1], steps)
        snr_results = []
        
        for ap in apertures:
            result = self.perform_photometry(ra, dec, ap)
            snr_results.append(result["snr"])
            logger.debug(f"孔径 {ap:.1f}\": SNR = {result['snr']:.1f}")

        # 绘制曲线
        plt.figure(figsize=(10, 6))
        plt.plot(apertures, snr_results, 'bo-', markersize=6)
        
        # 标记最佳孔径
        best_idx = np.argmax(snr_results)
        plt.annotate(f'Best: {apertures[best_idx]:.1f}"',
                    xy=(apertures[best_idx], snr_results[best_idx]),
                    xytext=(apertures[best_idx]+0.5, snr_results[best_idx]-5),
                    arrowprops=dict(facecolor='red', shrink=0.05))
        
        plt.xlabel('Aperture Radius (arcsec)')
        plt.ylabel('Signal-to-Noise Ratio')
        plt.title('SNR vs Aperture Size')
        plt.grid(True)
        
        if output_path:
            plt.savefig(output_path, dpi=150)
            logger.info(f"SNR曲线已保存至 {output_path}")
        else:
            plt.show()
        
        return apertures, snr_results

    
    
    def calculate_aperture_correction(self, ra, dec, 
                                    reference_aperture=5.0, 
                                    target_aperture=2.0,
                                    use_local_background=False):
        """
        计算孔径改正值
        :param reference_aperture: 参考孔径（arcsec），应足够大以包含全部流量
        :param target_aperture: 需要修正的目标孔径（arcsec）
        :return: 孔径改正值（星等差，需加至目标孔径测光结果）
        """
        # 使用参考孔径测光
        ref_result = self.perform_photometry(ra, dec, reference_aperture, use_local_background)
        
        # 使用目标孔径测光
        target_result = self.perform_photometry(ra, dec, target_aperture, use_local_background)
        
        # 计算孔径改正值
        ap_corr = ref_result['magnitude'] - target_result['magnitude']
        logger.info(f"孔径改正计算: {target_aperture}\" → {reference_aperture}\" | Δmag = {ap_corr:.3f}")
        
        return ap_corr

    def perform_photometry_with_correction(self, ra, dec, aperture_radius_arcsec, ap_corr=0.0, use_local_background=False):
        """
        执行带孔径改正的测光
        :param ap_corr: 孔径改正值（需加至原始星等）
        """
        # 原始测光结果
        result = self.perform_photometry(ra, dec, aperture_radius_arcsec, use_local_background)
        
        # 应用孔径改正
        if result['magnitude'] is not None:
            result['magnitude_corrected'] = result['magnitude'] + ap_corr
            logger.debug(f"应用孔径改正: {ap_corr:.3f} mag")
        else:
            result['magnitude_corrected'] = None
        
        return result

    def auto_aperture_correction(self, ra, dec, aperture_radius_arcsec=2.0,growth_range=(3.0, 8.0),
use_local_background=False):
        """
        自动孔径改正流程
        1. 生成生长曲线确定参考孔径
        2. 计算孔径改正值
        3. 返回修正后的测光结果
        """
        # 生成生长曲线确定稳定孔径
        growth_data = self.generate_mag_growth_curve(ra, dec, aperture_range=growth_range,steps=10)
        ref_ap = growth_data['optimal_aperture']
        logger.info(f"ref_ap:{ref_ap}")
        
        # 计算孔径改正
        ap_corr = self.calculate_aperture_correction(ra, dec,
                                                    reference_aperture=ref_ap,
                                                    target_aperture=aperture_radius_arcsec)
        
        # 应用改正后的测光
        return self.perform_photometry_with_correction(ra, dec, aperture_radius_arcsec, ap_corr=ap_corr,use_local_background=use_local_background)
if __name__ == "__main__":
    # 使用示例
    fitsfile = "/mnt/7b21f1e1-eb25-4cd5-bdb5-06d7d82fa253/Temp/force_photmetry/images/EP20250223T150435-25508/my_sc_tEP20250223T150435-25508_g_20250223153413_158_sciimg.fits"
    ra = 98.27392124999999   # 目标赤经
    dec = -22.444422777777778    # 目标赤纬
    (x, y), (height, width) = ra_dec_to_xy_and_size(fitsfile, ra, dec)
    partitioner = MatrixPartitioner(height, width)  # 将图像分成若干 gate
    logger.info(f"MatrixPartitioner:{partitioner}")
    range_info = partitioner.find_submatrix_with_range(x, y)
    logger.info(f"range_info:{range_info},(x, y):{(x, y)}")
    
    # 从原图像头部读取关键字，如曝光时间(EXPOSURE)、零点(ZPMAG) 等
    v_list = read_header_keys(fitsfile, ["exposure","ZPMAG","DATE","ZPMAG_G"])
    gate_find = range_info['submatrix_index'] - 1
    zero_point = float(v_list['ZPMAG_G'].split(',')[gate_find])
    logger.info(f"zero_point:{zero_point}")
    aperture_radius_arcsec = 1.98 * 1  # 测光孔径半径
    
    photometry = ForcedPhotometry(fitsfile, zero_point)
    result = photometry.perform_photometry(ra, dec, aperture_radius_arcsec, use_local_background=True)

    logger.info("Photometry Result:")
    logger.info(f"Total Flux: {result['total_flux']:.2f}")
    logger.info(f"Net Flux: {result['net_flux']:.2f}")
    logger.info(f"SNR: {result['snr']:.2f}")
    logger.info(f"Limiting Magnitude: {result['limiting_mag']:.2f}")
    logger.info(f"Forced Magnitude: {result['magnitude']:.2f}")
    
    
    # photometry.generate_light_curve(
    #     data_dir="/mnt/7b21f1e1-eb25-4cd5-bdb5-06d7d82fa253/Temp/force_photmetry/images/EP20250223T150435-25508",
    #     ra=ra,
    #     dec=dec,
    #     aperture_radius_arcsec=2.0,
    #     output_path="light_curve.png"
    # )
    
    # # 示例2: 生成SNR-孔径曲线
    # photometry.generate_snr_aperture_curve(
    #     ra=ra,
    #     dec=dec,
    #     aperture_range=(1.0, 5.0),
    #     steps=15,
    #     output_path="snr_curve.png"
    # )
    # growth_data = photometry.generate_mag_growth_curve(
    #     ra=ra,
    #     dec=dec,
    #     aperture_range=(0.5, 10),
    #     steps=30,
    #     output_path="mag_growth.png"
    # )
    # logger.info(f"最佳孔径: {growth_data['optimal_aperture']:.1f}\"")
    
    # # 方法1: 手动计算孔径改正
    # ap_corr = photometry.calculate_aperture_correction(ra, dec, reference_aperture=5.0,target_aperture=2.0)
    # corrected_result = photometry.perform_photometry_with_correction(ra, dec, 2.0, ap_corr=ap_corr)
    
    # 方法2: 全自动流程
    auto_corrected = photometry.auto_aperture_correction(ra, dec, 
                                                        aperture_radius_arcsec=2.2,
                                                        growth_range=(1.0, 10.0))
    